# $Id: rcm_sim_by_t.r,v 1.5 2006/03/28 18:09:24 jkatz Exp jkatz $
# rcm_sim_by_t.r
# by Jonathan N. Katz <jkatz@caltech.edu>
#
#
#############################################################
##                                                         ##
##      Load Libraries                                     ##
##                                                         ##
#############################################################
library(nlme)

#############################################################
##                                                         ##
##      User Setable Parameters                            ##
##                                                         ##
#############################################################
iter       <- 1000          # Number of simulation iterations 
n.groups   <- 20            # Number of units
b.mean     <- 5             # Mean of \beta_i
b.std      <- 1.8           # Standard deviation of \beta_i
reps.start <- 5             # Starting number of reps (T)
reps.end   <- 100           # Ending number of reps (T)
reps.step  <- 5             # Step size for reps

set.seed(666)          # Seed for RNG

#############################################################
##                                                         ##
##      Generate objects not changed during Design loop    ##
##                                                         ##
#############################################################

## Storage Arrays
design.length <-  length(seq(reps.start,reps.end,by=reps.step))

rmse.colnames  <-  c("n.reps", "b.mean","b.std", "b.group")

rmse.lm               <- array(NA, c(design.length,4))
colnames(rmse.lm)     <- rmse.colnames

rmse.lmList           <- array(NA, c(design.length,4))
colnames(rmse.lmList) <-  rmse.colnames

rmse.lmeML            <- array(NA, c(design.length,4))
colnames(rmse.lmeML)  <-  rmse.colnames

rmse.lmeREML          <- array(NA, c(design.length,4))
colnames(rmse.lmeREML)<-  rmse.colnames

rmse.lmeFGLS          <- array(NA, c(design.length,4))
colnames(rmse.lmeFGLS)<-  rmse.colnames

### Generate \beta_i's for all runs
b.group.all <-  array(rnorm(n.groups*iter, mean=b.mean, sd=b.std),
                      c(iter,n.groups))

############################################################
##                                                         ##
##      Design Loop                                        ##
##                                                         ##
#############################################################

j <-  1
for (n.reps in seq(reps.start,reps.end,by=reps.step)) {
  cat("\n", "Expermiment, t= ", n.reps,"\n" )

#############################################################
##                                                         ##
##      Generate objects not changed during MC looping     ##
##                                                         ##
#############################################################
  
### Useful fixed values
  n         <- n.groups*n.reps
  group.id  <- rep (1:n.groups, each=n.reps)
 
  
#### Arrays to store results
  
  results.colnames          <-  c("b.mean","b.std",paste("b",1:n.groups,sep=""))
  
  results.lm                <- array(NA,c(iter,n.groups+2))
  colnames(results.lm)      <- results.colnames
  
  results.lmList            <- array(NA,c(iter,n.groups+2))  
  colnames(results.lmList)  <- results.colnames
  
 
  results.lmeML             <- array(NA,c(iter,n.groups+2))
  colnames(results.lmeML)   <- results.colnames
  
  results.lmeREML           <- array(NA,c(iter,n.groups+2))
  colnames(results.lmeREML) <- results.colnames

  results.lmeFGLS           <- array(NA,c(iter,n.groups+2))
  colnames(results.lmeFGLS) <- results.colnames

#############################################################
##                                                         ##
##      Monte Carlo Loop                                   ##
##                                                         ##
#############################################################

  for (i in 1:iter) {
    
### Where are we in the loop
    cat(i, " ")    

    
### Generate X's
    x <- rnorm(n,mean=1,sd=0.10)

### Generate Y's
    b.group  <- b.group.all[i,]
    y        <- b.group[group.id]*x + rnorm(n, mean=0, sd=1)
    sim.data <- data.frame(y=y,x=x,group.id=as.factor(group.id))
    
### Pooled OLS
    sim.lm <-  lm(y ~  x, data=sim.data)
    results.lm[i,] <- unlist(c(coef(sim.lm)[2],NA,
                               rep(coef(sim.lm)[2],n.groups)))
    
### Unit by Unit OLS              
    sim.lmList <-  lm(y ~ group.id/(x-1), data=sim.data)
    results.lmList[i,] <-  unlist(c(NA,NA,
                                    coef(sim.lmList)[(n.groups+1):(2*n.groups)]
                                    ))

### RCM via ML
    sim.lmeML <-  lme(y ~ x, random= ~ 0 + x  | group.id, data=sim.data,
                      method="ML")
    results.lmeML[i,] <-  unlist(c(fixef(sim.lmeML)[2],
                                 sqrt(as.numeric(getVarCov(sim.lmeML))),
                                   coef(sim.lmeML)[,2]))
    
### RCM via REML
    sim.lmeREML <-  lme(y ~ x ,random= ~ 0 + x | group.id, data=sim.data,
                        method="REML")
    results.lmeREML[i,] <-  unlist(c(fixef(sim.lmeREML)[2],
                                     sqrt(as.numeric(getVarCov(sim.lmeREML))),
                                     coef(sim.lmeREML)[,2]))

### RCM via FGLS
    b.i      <-  coef(sim.lmList)[(n.groups+1):(2*n.groups)]
    v.i      <-  diag(vcov(sim.lmList))[(n.groups+1):(2*n.groups)]
#    g      <-  var(b.i) - mean(v.i)
    g        <-  var(b.i)
    w        <- (1/sum(1/(g + v.i)))*(1/(g + v.i))
    b.fgls   <- sum(w*b.i)
    a.i      <-  (1/((1/v.i) + (1/g)))*(1/g)
    b.i.fgls <-  a.i*b.fgls + (1 - a.i)*b.i

    results.lmeFGLS[i,] <-  unlist(c(b.fgls,g,b.i.fgls))

  }
  
  
#### Calculate RMSE
  
# This cose creates and expands the true parameters to a matrix the
# same size as the results.

  truth <-  c(b.mean,b.std)
  truth <-  t(array(truth,c(2,iter)))
  truth <-  cbind(truth,b.group.all)
  
  mse.lm           <- apply((results.lm-truth)^2,2,mean)
  rmse.lm[j,]      <- unlist(c(n.reps, sqrt(mse.lm[1:2]),
                               sqrt(mean(mse.lm[3:length(mse.lm)]))))
  
  mse.lmList       <- apply((results.lmList-truth)^2,2,mean)
  rmse.lmList[j,]  <- unlist(c(n.reps, sqrt(mse.lmList[1:2]),
                               sqrt(mean(mse.lmList[3:length(mse.lmList)]))))
  
  mse.lmeML        <- apply((results.lmeML-truth)^2,2,mean)
  rmse.lmeML[j,]   <- unlist(c(n.reps, sqrt(mse.lmeML[1:2]),
                               sqrt(mean(mse.lmeML[3:length(mse.lmeML)]))))
  
  mse.lmeREML      <- apply((results.lmeREML-truth)^2,2,mean)
  rmse.lmeREML[j,] <- unlist(c(n.reps, sqrt(mse.lmeREML[1:2]),
                               sqrt(mean(mse.lmeREML[3:length(mse.lmeREML)]))))

  mse.lmeFGLS      <- apply((results.lmeFGLS-truth)^2,2,mean)
  rmse.lmeFGLS[j,] <- unlist(c(n.reps, sqrt(mse.lmeFGLS[1:2]),
                               sqrt(mean(mse.lmeFGLS[3:length(mse.lmeFGLS)]))))

  save(rmse.lm, rmse.lmList, rmse.lmeML, rmse.lmeREML, rmse.lmeFGLS,
       file="rmse_sim_by_T.Rdata", ascii=TRUE)

j <-  j+1
}


# In order to do the FGLS estimate, extract the vcov matrix using
# lapply(test.lis, vcov)

# Trick to get the se from the lme() 
# summary(sim.lmeML)$tTable[2]
 #do something like save(rmse.lm, rmse.,,, file="rsme.sim.Rdata")
